import arrow
from utils.date_utils import dynamic_format


def build_daily_query(timestamp: float, imei: str, days_before: int = 0,
                      days_after: int = 0) -> dict:
    """生成指定天数范围内的数据查询ES Query

    Args:
        timestamp (float): 时间戳
        imei (str): 设备IMEI号
        days_before (int, optional): 查询前几天的数据，默认为0
        days_after (int, optional): 查询后几天的数据，默认为0

    Returns:
        dict: Elasticsearch查询体
    """
    # 转换时间戳为arrow对象并设置时区
    dt_with_tz = arrow.get(timestamp, tzinfo='Asia/Shanghai')

    # 计算查询起始日期 (使用arrow的shift方法)
    start_date = dt_with_tz.shift(days=-days_before).floor('day')

    # 计算查询结束日期
    end_date = dt_with_tz.shift(days=days_after).ceil('day').shift(seconds=-1)

    # 生成ES兼容的ISO格式，转换到UTC
    gte = dynamic_format(start_date.format('YYYY-MM-DDTHH:mm:ss.SSSZ'))
    lte = dynamic_format(end_date.format('YYYY-MM-DDTHH:mm:ss.SSSZ'))

    # 构建查询体
    return {
        "query": {
            "bool": {
                "must": [
                    {"terms": {"imei": [imei]}}
                ],
                "filter": [
                    {"range": {"dt": {"gte": gte, "lte": lte}}}
                ],
                "must_not": [
                    {"term": {"validState": {"value": 0}}}
                ]
            }
        },
        "sort": [
            {"dt": {"order": "asc"}}
        ]
    }